package org.toycode.solr.extractor.impl;

import com.chenlb.mmseg4j.Dictionary;
import com.chenlb.mmseg4j.example.Complex;
import java.io.IOException;
import java.io.Writer;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.List;
import java.util.logging.Level;
import java.util.logging.Logger;
import org.apache.solr.common.SolrInputDocument;
import org.pdfbox.pdmodel.PDDocument;
import org.pdfbox.util.PDFTextStripper;
import org.toycode.solr.config.Prop;
import org.toycode.solr.extractor.PropertyExtractor;

/**
 * This implementations of the interface {@linkplain PropertyExtractor}
 * provides the content extraction for the type <em>pdf</em> and
 * <em>pdfs</em> in config.xml.
 *
 * @author <a href="mailto:DL88250@gmail.com">Liang Ding</a>
 * @version 1.0.0.1, Aug 26, 2009
 */
final public class PdfExtractor implements PropertyExtractor {

    private final List<String> EXTRACTABLE_TYPES = Arrays.asList("pdf",
                                                                 "pdfs");

    @SuppressWarnings("unchecked")
    public void process(final Prop prop, final Object value,
                        final SolrInputDocument solrInputDocument) {
        final String type = prop.type;
        if (type.equals("pdf")) {
            final String content = readPdf((String) value);
            solrInputDocument.addField("pdf_content", content);
        } else if (type.equals("pdfs")) {
            final List<String> contents = new ArrayList<String>();
            for (final String path : ((Collection<String>) value)) {
                contents.add(readPdf(path));
            }
            solrInputDocument.addField("pdf_content", contents);
        }
    }

    private static String readPdf(final String path) {
        boolean sort = false;
        String pdfFile = path;
        String encoding = "UTF-8";
        int startPage = 1;
        int endPage = Integer.MAX_VALUE;
        Writer output = null;
        PDDocument document = null;
        try {
            document = PDDocument.load(pdfFile);
            PDFTextStripper stripper = null;
            stripper = new PDFTextStripper();
            stripper.setSortByPosition(sort);
            stripper.setStartPage(startPage);
            stripper.setEndPage(endPage);
            return stripper.getText(document);
        } catch (final IOException ex) {
            Logger.getLogger(PdfExtractor.class.getName()).
                    log(Level.SEVERE, null, ex);
        } finally {
            if (output != null) {
                try {
                    output.close();
                } catch (final IOException ex) {
                    Logger.getLogger(PdfExtractor.class.getName()).
                            log(Level.SEVERE, null, ex);
                }
            }
            if (document != null) {
                try {
                    document.close();
                } catch (final IOException ex) {
                    Logger.getLogger(PdfExtractor.class.getName()).
                            log(Level.SEVERE, null, ex);
                }
            }
        }

        return null;
    }

    public static void main(String[] args) throws IOException {
        final String text = readPdf("src/test/resources/sample1.pdf");
        final String segWords = new Complex(
                new Dictionary("src/test/resources/solr/dic")).segWords("目录测试",
                                                                        " | ");
        System.out.println(segWords);
    }

    public List<String> getExtractableTypes() {
        return EXTRACTABLE_TYPES;
    }
}
